package org.idea.qiyu.framework.apm.agent.core.servlet;

import javassist.*;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;
import java.util.UUID;

/**
 * @Author linhao
 * @Date created in 11:41 上午 2021/7/18
 */
public class ServletAgent {

    private static final String SERVLET_CLASS = "javax/servlet/http/HttpServlet";



    public static void premain(String args, Instrumentation instrumentation) {
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (loader == null || className == null) {
                    return null;
                }
                if (!SERVLET_CLASS.equals(className)) {
                    return null;
                }
                return buildAgentClassByte(loader,className.replace("/", "."));
            }
        },true);
    }

    public static byte[] buildAgentClassByte(ClassLoader classLoader,String className){
            try {
                System.out.println("======");
                ClassPool classPool = new ClassPool();
                classPool.appendSystemPath();
                CtClass ctClass = classPool.get(className);
                CtMethod ctMethod = ctClass.getDeclaredMethod("service",
                        ClassPool.getDefault().get(new String[]{
                                "javax.servlet.http.HttpServletRequest", "javax.servlet.http.HttpServletResponse"
                        }));
                CtMethod copyMethod = CtNewMethod.copy(ctMethod, ctClass, new ClassMap());
                String newMethodName = ctMethod.getName() + "$Agent";
                ctMethod.setName(newMethodName);
                System.out.println("!!!!");
                if ("void".equals(ctMethod.getReturnType())) {
                    System.out.println("void");
                    copyMethod.setBody(" {\n" +
                            "                            Object object = org.idea.qiyu.framework.apm.agent.core.servlet.ServletAgent.begin($$);\n" +
                            "                            try {\n" +
                            "                                " + newMethodName + "($$);\n" +
                            "                            }finally {\n" +
                            "                                org.idea.qiyu.framework.apm.agent.core.servlet.ServletAgent.end(object);\n" +
                            "                            }\n" +
                            "                        }");
                } else {
                    System.out.println("return");
                    copyMethod.setBody(" {\n" +
                            "                            Object object = org.idea.qiyu.framework.apm.agent.core.servlet.ServletAgent.begin($$);\n" +
                            "                            try {\n" +
                            "                                return " + newMethodName + "($$);\n" +
                            "                            }finally {\n" +
                            "                                org.idea.qiyu.framework.apm.agent.core.servlet.ServletAgent.end(object);\n" +
                            "                            }\n" +
                            "                        }");
                }
                ctClass.addMethod(copyMethod);
                return ctClass.toBytecode();
            } catch (Exception e) {
                System.out.println("error , e:"+e.getCause());
                e.printStackTrace();
            }
            return null;
    }

    public static ServletStatisticsObject begin(HttpServletRequest request, HttpServletResponse response) {
        ServletStatisticsObject servletStatisticsObject = new ServletStatisticsObject();
        servletStatisticsObject.setUrl(request.getRequestURI());
        servletStatisticsObject.setBegin(System.currentTimeMillis());
        servletStatisticsObject.setCookies(request.getCookies());
        servletStatisticsObject.setParams(request.getParameterMap());
        servletStatisticsObject.setTraceId(UUID.randomUUID().toString());
        return servletStatisticsObject;
    }

    public static ServletStatisticsObject end(Object object) {
        ServletStatisticsObject servletStatisticsObject = (ServletStatisticsObject) object;
        servletStatisticsObject.setEnd(System.currentTimeMillis());
        System.out.println(servletStatisticsObject.toString());
        return servletStatisticsObject;
    }

}
